modelopt/torch/puzzletron/tools/logger.py-164-168 (1)

164-168: ⚠️ Potential issue | 🟡 Minor

Docstring says "rank -1 in node 0" but implementation may not match.

The lprint docstring says "Last rank prints only (rank -1 in node 0)" but the implementation checks local_rank != world_size - 1. If the intent is the last rank globally, this logic appears incorrect. If the intent is the last local rank within a node, the comparison should be against local world size, not global world size.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@modelopt/torch/puzzletron/tools/logger.py` around lines 164 - 168, The
docstring for lprint mismatches its rank-check logic; update either the
docstring or the implementation in lprint to match the intended behavior: if you
want "last rank in node 0" (last local rank) change the check to compare
local_rank against local_world_size - 1 (use the local world size variable used
elsewhere in logger or from torch/distributed) before calling logger.dist_log,
or if you intended the global last rank adjust the docstring to "last global
rank" and ensure the check compares local_rank against world_size - 1
accordingly; reference the lprint function and logger.dist_log to locate where
to change the comparison and/or docstring.

modelopt/torch/puzzletron/tools/logger.py-62-66 (1)

62-66: ⚠️ Potential issue | 🟡 Minor

Error message is inconsistent with valid choices.

The error message on line 65 lists ['all', 'main', 'local_main'] but 'last' is also a valid choice per line 62.

🐛 Proposed fix

         if ranks not in ["all", "main", "local_main", "last"]:
             raise NotImplementedError(
                 f"Could not broadcast msg {msg} - "
-                f"ranks parameters choices are ['all', 'main', 'local_main']. Got {ranks}"
+                f"ranks parameters choices are ['all', 'main', 'local_main', 'last']. Got {ranks}"
             )

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@modelopt/torch/puzzletron/tools/logger.py` around lines 62 - 66, The
NotImplementedError message for the ranks check mentions the wrong set of valid
choices; update the f-string in the ranks validation (the raise
NotImplementedError block that references the variable ranks and msg) to include
'last' in the displayed list (i.e., ["all", "main", "local_main", "last"]) so
the error accurately reflects the valid options.

examples/pruning/README.md-10-10 (1)

10-10: ⚠️ Potential issue | 🟡 Minor

Fix compound adjective hyphenation in docs.

At Line 10, use “MIP-based” (not “MIP based”) for correct technical phrasing.

✏️ Proposed edit

-1. [Puzzletron](../puzzletron/README.md): An advanced pruning method by NVIDIA using Mixed Integer Programming (MIP) based NAS search algorithm.
+1. [Puzzletron](../puzzletron/README.md): An advanced pruning method by NVIDIA using a Mixed Integer Programming (MIP)-based NAS search algorithm.

As per coding guidelines, "**/*.{py,cpp,h,hpp,md}: All code (Python, C++, Markdown, etc.) must adhere to coding standards as checked by pre-commit hooks".

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@examples/pruning/README.md` at line 10, Update the compound adjective "MIP
based" in the Puzzletron README entry to the hyphenated form "MIP-based" so the
phrase reads "MIP-based NAS search algorithm"; locate the text string
"Puzzletron" / the line containing "MIP based NAS search algorithm" and replace
it with "MIP-based" to comply with the Markdown/style pre-commit checks.

examples/puzzletron/configs/gptoss-20b_remove_experts_memory/gptoss-20b_remove_experts_memory.yaml-14-17 (1)

14-17: ⚠️ Potential issue | 🟡 Minor

Fix the mismatch between target_memory and its comment.

16_000 MiB is about 15.6 GiB; 45 GiB would be 46_080 MiB. One of these is wrong, and the current ambiguity can make the example optimize for the wrong budget.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In
`@examples/puzzletron/configs/gptoss-20b_remove_experts_memory/gptoss-20b_remove_experts_memory.yaml`
around lines 14 - 17, The mip.human_constraints.target_memory value
(target_memory: 16_000) and its comment (“45 GiB”) disagree; update one to match
the other: either set target_memory to 46_080 (MiB) if you mean 45 GiB, or
change the comment to match 16_000 MiB (~15.6 GiB). Locate the target_memory key
under mip -> human_constraints and make the numeric value and the GiB comment
consistent.

examples/puzzletron/configs/nemotron-nano-12b-v2/nemotron_nano_12b_v2_pruneffn_memory.yaml-14-17 (1)

14-17: ⚠️ Potential issue | 🟡 Minor

Align the MiB value with the GiB comment.

90_000 MiB is about 87.9 GiB, not 90 GiB. If the intended cap is 90 GiB, this should be 92_160; otherwise update the comment so the search target is unambiguous.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In
`@examples/puzzletron/configs/nemotron-nano-12b-v2/nemotron_nano_12b_v2_pruneffn_memory.yaml`
around lines 14 - 17, The target_memory under mip.human_constraints is
inconsistent with the "90 GiB" comment; update the value of target_memory to
92_160 (MiB) to represent 90 GiB, or alternatively change the comment to the
actual GiB represented by 90_000 MiB (~87.9 GiB); edit the target_memory key in
this YAML (mip -> human_constraints -> target_memory) accordingly to keep value
and comment aligned.

examples/puzzletron/configs/llama-3_2-3B_pruneffn_memory/llama-3_2-3B_pruneffn_memory.yaml-17-17 (1)

17-17: ⚠️ Potential issue | 🟡 Minor

Memory unit comment is inaccurate.

Line 17 uses 45_000 MiB, which is ~43.95 GiB, not 45 GiB. Please update either the value or the comment to avoid misconfigured constraints.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In
`@examples/puzzletron/configs/llama-3_2-3B_pruneffn_memory/llama-3_2-3B_pruneffn_memory.yaml`
at line 17, The comment next to the target_memory value is incorrect:
target_memory: 45_000 is in MiB (~43.95 GiB), not 45 GiB; either set
target_memory to 46_080 (45 * 1024) to represent 45 GiB exactly or update the
inline comment to the correct GiB approximation (≈43.95 GiB); locate the
target_memory entry and apply one of these fixes so the value and its comment
match.

modelopt/torch/puzzletron/anymodel/models/llama/llama_converter.py-36-52 (1)

36-52: ⚠️ Potential issue | 🟡 Minor

Return type annotation mismatch: method declares List[BlockConfig] but returns List[dict].

The method signature declares a return type of List[BlockConfig], but line 50 calls .to_dict() on each BlockConfig instance, causing the method to actually return List[dict]. Update the return type annotation to match the actual behavior.

🔧 Suggested fix

     `@staticmethod`
-    def create_block_configs_from_main_config(config: LlamaConfig) -> List[BlockConfig]:
+    def create_block_configs_from_main_config(config: LlamaConfig) -> List[dict]:
         """Create uniform block configs for all Llama layers.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@modelopt/torch/puzzletron/anymodel/models/llama/llama_converter.py` around
lines 36 - 52, The function create_block_configs_from_main_config currently
builds BlockConfig objects but calls .to_dict(), so its declared return type
List[BlockConfig] is incorrect; update the signature to return List[dict] (or
List[MutableMapping] if you prefer an abstract type) to match the actual
returned value, and ensure the typing import (List) is present/adjusted
accordingly; reference create_block_configs_from_main_config, LlamaConfig,
BlockConfig, and the .to_dict() call when making the change.

modelopt/torch/puzzletron/anymodel/converter/converter_factory.py-43-47 (1)

43-47: ⚠️ Potential issue | 🟡 Minor

Handle potential None from inspect.getmodule().

inspect.getmodule(ref) can return None if the module cannot be determined, which would cause an AttributeError when accessing .__name__.

🛡️ Proposed fix

-                raise KeyError(
-                    f"Could not register `{cls_name}`: {cls_type}, "
-                    f"`{cls_name}` is already registered and points to "
-                    f"`{inspect.getmodule(ref).__name__}.{ref.__name__}`"
-                )
+                module = inspect.getmodule(ref)
+                module_name = module.__name__ if module else "<unknown>"
+                raise KeyError(
+                    f"Could not register `{cls_name}`: {cls_type}, "
+                    f"`{cls_name}` is already registered and points to "
+                    f"`{module_name}.{ref.__name__}`"
+                )

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@modelopt/torch/puzzletron/anymodel/converter/converter_factory.py` around
lines 43 - 47, The KeyError message in converter_factory.py assumes
inspect.getmodule(ref) always returns a module; update the error construction in
the registration failure branch (where cls_name, cls_type and ref are used) to
guard against None from inspect.getmodule(ref) by using a safe fallback (e.g.,
module = inspect.getmodule(ref); module_name = module.__name__ if module is not
None else getattr(ref, "__module__", "unknown")) and use module_name rather than
inspect.getmodule(ref).__name__; ensure ref.__name__ (or getattr(ref,
"__name__", repr(ref))) is similarly accessed safely so the f-string does not
raise AttributeError.

modelopt/torch/puzzletron/puzzletron.py-60-74 (1)

60-74: ⚠️ Potential issue | 🟡 Minor

Step numbering is inconsistent.

Comments jump from "Step 2" to "Step 4", skipping "Step 3". This appears to be a typo or leftover from a refactor. Consider renumbering for clarity:

📝 Proposed fix

     # Step 2: pruning_ckpts (single process)
     if dist.is_master():
         pruning_ckpts.launch_prune_ckpt(hydra_cfg)
     dist.barrier()

-    # Step 4: build_library_and_stats (single process)
+    # Step 3: build_library_and_stats (single process)
     if dist.is_master():
         build_library_and_stats.launch_build_library_and_stats(hydra_cfg)
     dist.barrier()

-    # Step 5: calc_one_block_scores (distributed processing)
+    # Step 4: calc_one_block_scores (distributed processing)
     scoring.launch_scoring(hydra_cfg)

-    # Step 6: mip_and_realize_models (distributed processing)
+    # Step 5: mip_and_realize_models (distributed processing)
     mip_and_realize_models.launch_mip_and_realize_model(hydra_cfg)

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@modelopt/torch/puzzletron/puzzletron.py` around lines 60 - 74, The step
comments are misnumbered (jump from "Step 2" to "Step 4"); update the inline
step comments in puzzletron.py to use a consistent sequence (e.g., Step 2 → Step
3 → Step 4 etc.) around the blocks that call pruning_ckpts.launch_prune_ckpt,
build_library_and_stats.launch_build_library_and_stats, scoring.launch_scoring,
and mip_and_realize_models.launch_mip_and_realize_model so the step numbers
reflect the actual order of execution.

examples/puzzletron/configs/nemotron-nano-12b-v2/pruning/pruning_defaults.yaml-20-31 (1)

20-31: ⚠️ Potential issue | 🟡 Minor

Remove unused ffn_list definition or align with intermediate_size_list usage.

The file serves as a defaults template for Hydra configuration, but there's an inconsistency: ffn_list is defined on line 21 with no value, yet the actual FFN pruning configuration (ffn_pruning.yaml) populates and uses intermediate_size_list instead. Either rename ffn_list to intermediate_size_list for consistency, or remove the unused ffn_list definition.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In
`@examples/puzzletron/configs/nemotron-nano-12b-v2/pruning/pruning_defaults.yaml`
around lines 20 - 31, The defaults file defines an unused key ffn_list while the
FFN pruning logic uses intermediate_size_list (see ffn_pruning.yaml); fix by
either removing the ffn_list entry entirely or renaming it to
intermediate_size_list and ensuring the init mode key (mlp_init_mode) matches
the expected name used by ffn_pruning.yaml so the Hydra defaults align with the
pruning config.

modelopt/torch/puzzletron/tools/hydra_utils.py-29-36 (1)

29-36: ⚠️ Potential issue | 🟡 Minor

Validate block and mbs before division in warmup_steps.

block=0 or mbs=0 causes a runtime failure. Add explicit positive-value checks and fail early with a clear error.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@modelopt/torch/puzzletron/tools/hydra_utils.py` around lines 29 - 36, The
warmup_steps function currently divides by block and mbs which will raise on
zero; update warmup_steps to validate that block and mbs are positive integers
(e.g., block > 0 and mbs > 0) at the top of the function and raise a clear
ValueError if not; keep the existing logic (compute steps =
(tokens//block)//mbs, compute w = pct * steps, return max(1, round(w))) but fail
fast with an explanatory error mentioning the parameter names block and mbs so
callers get a clear message when invalid values are passed.

modelopt/torch/puzzletron/activation_scoring/activation_hooks/utils.py-46-47 (1)

46-47: ⚠️ Potential issue | 🟡 Minor

Avoid mutating caller-owned activation_hooks_kwargs.

This function modifies the input dict in place (["model"] = model), which can leak state across calls. Copy it first before enrichment.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@modelopt/torch/puzzletron/activation_scoring/activation_hooks/utils.py`
around lines 46 - 47, The code mutates the caller-owned dict
activation_hooks_kwargs by setting activation_hooks_kwargs["model"] = model;
instead, make a shallow copy of activation_hooks_kwargs (e.g., new_kwargs =
activation_hooks_kwargs.copy() or dict(activation_hooks_kwargs)), add the
"model" key to the copy, and use/return that copy so the original dict isn’t
mutated; update any return/usage to use the new_kwargs variable.

examples/puzzletron/configs/llama-3_1-8B_pruneffn_memory/Llama-3_1-8B.yaml-13-13 (1)

13-13: ⚠️ Potential issue | 🟡 Minor

Minor typo in comment.

ppath should be path.

Proposed fix

-dataset_path: ??? # ppath to Nemotron-Post-Training-Dataset-v2
+dataset_path: ??? # path to Nemotron-Post-Training-Dataset-v2

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@examples/puzzletron/configs/llama-3_1-8B_pruneffn_memory/Llama-3_1-8B.yaml`
at line 13, Fix the typo in the YAML comment for the dataset key: change the
comment after the dataset_path entry so it reads "path to
Nemotron-Post-Training-Dataset-v2" instead of "ppath to
Nemotron-Post-Training-Dataset-v2" so the comment for the dataset_path key is
correct and clear.

modelopt/torch/puzzletron/tools/validate_model.py-138-140 (1)

138-140: ⚠️ Potential issue | 🟡 Minor

Potential issue: validation_full_iters could be zero.

If args.eval_samples is less than args.micro_batch_size, the integer division results in validation_full_iters = 0. This could cause unexpected behavior in downstream code that expects at least one iteration.

🛡️ Proposed fix: Add validation or use ceiling division

     validation_full_iters = (
-        args.eval_samples // args.micro_batch_size
+        max(1, args.eval_samples // args.micro_batch_size)
     )  # model pipeline, single data rank

Or add an assertion:

     validation_full_iters = (
         args.eval_samples // args.micro_batch_size
     )  # model pipeline, single data rank
+    assert validation_full_iters > 0, f"eval_samples ({args.eval_samples}) must be >= micro_batch_size ({args.micro_batch_size})"

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@modelopt/torch/puzzletron/tools/validate_model.py` around lines 138 - 140,
The calculation of validation_full_iters can yield zero when args.eval_samples <
args.micro_batch_size; update the logic around validation_full_iters in
validate_model.py to ensure at least one validation iteration by either
validating inputs (raise/exit if args.eval_samples < args.micro_batch_size) or
computing a ceiling division and then taking max(1, ...); reference the variable
validation_full_iters and the CLI args args.eval_samples and
args.micro_batch_size when making the change so downstream loops that expect >=1
iteration do not break.

modelopt/torch/puzzletron/subblock_stats/calc_subblock_stats.py-316-321 (1)

316-321: ⚠️ Potential issue | 🟡 Minor

Actually emit the missing-MoE-stats warning.

Warning(...) only instantiates an object; it does not notify the caller. Right now a missing moe_stats.json silently changes active_params from MoE-aware counts to total params. Use warnings.warn(...) or mprint(...) here.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@modelopt/torch/puzzletron/subblock_stats/calc_subblock_stats.py` around lines
316 - 321, The code creates a pathlib object moe_stats_file and on missing file
only instantiates Warning(...) which does nothing; update the block around
moe_stats_file (using variables moe_stats_file, moe_stats_filename,
master_puzzle_dir and the downstream active_params logic) to actually emit a
warning by calling warnings.warn(...) (or mprint(...)) with a clear message like
f"MOE stats file {moe_stats_file} does not exist, can't calculate num active
params", and add `import warnings` at the top if using warnings.warn; keep
setting moe_stats_file = None so the fallback to total params remains.

modelopt/torch/puzzletron/decilm/deci_lm_hf_code/block_config.py (3)

76-82: Consider deleting frame reference to avoid reference cycles.

Per Python documentation, holding references to frame objects can prevent garbage collection. Delete the frame reference when done.

♻️ Proposed fix

 def _is_called_from_post_init() -> bool:
     frame = inspect.currentframe()
-    while frame:
-        if frame.f_code.co_name == "__post_init__":
-            return True
-        frame = frame.f_back
-    return False
+    try:
+        while frame:
+            if frame.f_code.co_name == "__post_init__":
+                return True
+            frame = frame.f_back
+        return False
+    finally:
+        del frame

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@modelopt/torch/puzzletron/decilm/deci_lm_hf_code/block_config.py` around
lines 76 - 82, The function _is_called_from_post_init currently retains a
reference to the inspected frame object which can create reference cycles;
update the function to ensure the frame reference is cleared when done (e.g.,
use a try/finally or explicitly del frame before returning) so the frame object
can be garbage-collected; keep the existing traversal logic (checking
frame.f_code.co_name == "__post_init__") but ensure you delete the frame
variable at the end of the function and also delete any temporary frame
references used inside the loop.

---

89-93: Bare except: catches too much.

This catches SystemExit, KeyboardInterrupt, and other exceptions that shouldn't be silently swallowed. Use except Exception: to catch only regular exceptions.

♻️ Proposed fix

     try:
         _get_dataclass_type(tp)
         return True
-    except:
+    except Exception:
         return False

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@modelopt/torch/puzzletron/decilm/deci_lm_hf_code/block_config.py` around
lines 89 - 93, The current try/except around the call to _get_dataclass_type
uses a bare except which swallows SystemExit/KeyboardInterrupt; change the bare
except to catch only regular exceptions (use except Exception:) in the block
that tries _get_dataclass_type and returns True/False so only normal errors are
caught while allowing critical exceptions to propagate.

---

96-96: Incorrect return type annotation.

dataclass is a decorator, not a type. The return type should be Type[Any] or similar.

♻️ Proposed fix

-def _get_dataclass_type(tp: Type) -> dataclass:
+def _get_dataclass_type(tp: Type) -> Type[Any]:

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@modelopt/torch/puzzletron/decilm/deci_lm_hf_code/block_config.py` at line 96,
The return type annotation on _get_dataclass_type is wrong: `dataclass` is a
decorator, not a type; change the signature to return a type object (e.g. use ->
Type[Any] or a generic -> Type[T]) and update imports as needed (from typing
import Any, Type or TypeVar T) so the function signature becomes something like
def _get_dataclass_type(tp: Type) -> Type[Any]:; ensure any callers/types align
with the new typing.

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

modelopt/torch/puzzletron/subblock_stats/calc_subblock_stats.py (1)

404-415: ⚠️ Potential issue | 🟠 Major

Return format mismatch in fallback loader (existing issue).

The _load_subblock_configs_from_subblock_library function returns raw AttentionConfig/FFNConfig objects, but calculate_subblock_stats (lines 151-152) expects items with ["subblock_config"] and ["parent_layer_indices"] keys. This will cause a runtime error when replacement_library.json is absent.

🐛 Proposed fix

 def _load_subblock_configs_from_subblock_library(master_puzzle_dir: Path) -> list[SubblockConfig]:
     subblocks_df = pd.read_json(master_puzzle_dir / "subblock_library.json")
     subblocks_df["attention_config"] = subblocks_df["attention_config"].apply(
         partial(_dataclass_from_dict, cls=AttentionConfig)
     )
     subblocks_df["ffn_config"] = subblocks_df["ffn_config"].apply(
         partial(_dataclass_from_dict, cls=FFNConfig)
     )
     attention_configs = subblocks_df["attention_config"].dropna().drop_duplicates().tolist()
     ffn_configs = subblocks_df["ffn_config"].dropna().drop_duplicates().tolist()
-    subblock_configs = attention_configs + ffn_configs
+    subblock_configs = [
+        immutabledict({"subblock_config": cfg, "parent_layer_indices": tuple([-1])})
+        for cfg in attention_configs + ffn_configs
+    ]
     return subblock_configs

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@modelopt/torch/puzzletron/subblock_stats/calc_subblock_stats.py` around lines
404 - 415, The fallback loader _load_subblock_configs_from_subblock_library
currently returns raw AttentionConfig/FFNConfig objects but
calculate_subblock_stats expects each item to be a mapping with keys
"subblock_config" and "parent_layer_indices"; change
_load_subblock_configs_from_subblock_library to wrap each config into a dict
like {"subblock_config": <config>, "parent_layer_indices": []} (use an empty
list or appropriate default for parent_layer_indices) for both attention_configs
and ffn_configs so the combined subblock_configs list matches the shape expected
by calculate_subblock_stats.

modelopt/torch/puzzletron/mip/run_puzzle.py (1)

396-424: ⚠️ Potential issue | 🟠 Major

CLI validation crashes on argparse.Namespace (existing issue).

The checks like "subblock_stats_args" not in args (line 408) use in operator on argparse.Namespace, which is not iterable and will raise TypeError. The past review comment already flagged this issue.

🐛 Proposed fix

 def _assert_valid_config(args, puzzle_profile):
     required_args = (
         "subblock_stats_path",
         "objective",
         "output_path",
     )
-    missing_args = [arg for arg in required_args if arg not in args or getattr(args, arg) is None]
+    missing_args = [arg for arg in required_args if getattr(args, arg, None) is None]
     if missing_args:
         mprint(f"error: The following arguments are required: {', '.join(missing_args)}")
         sys.exit(1)

     # Make sure we have specified subblock_stats_args
-    if "subblock_stats_args" not in args and "subblock_stats_args" not in puzzle_profile:
+    if getattr(args, "subblock_stats_args", None) is None and "subblock_stats_args" not in puzzle_profile:
         mprint(
             "error: Must specify `subblock_stats_arrs` in either puzzle_profile or as a commandline arg."
         )
         sys.exit(1)

     # Make sure we have specified constraints
     if (
-        "mip_constraints" not in args
-        and "human_constraints" not in args
+        getattr(args, "mip_constraints", None) is None
+        and getattr(args, "human_constraints", None) is None
         and "mip_constraints" not in puzzle_profile
         and "human_constraints" not in puzzle_profile
     ):

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@modelopt/torch/puzzletron/mip/run_puzzle.py` around lines 396 - 424, In
_assert_valid_config, the checks like `"subblock_stats_args" not in args` and
similar use the in-operator on argparse.Namespace (causing TypeError); change
those presence checks to use hasattr(args, "subblock_stats_args") or
getattr(args, "subblock_stats_args") is not None for each CLI arg check, while
keeping the existing membership checks against puzzle_profile (which is a dict);
update the two conditional blocks that validate subblock_stats_args and the
constraints (mip_constraints, human_constraints) to use hasattr/getattr on args
and keep the puzzle_profile checks as-is so validation works for both Namespace
and profile.

modelopt/torch/puzzletron/anymodel/models/qwen3/qwen3_converter.py (1)

33-33: Return type annotation mismatch with actual return value.

The signature declares List[BlockConfig] but the method returns List[dict] after calling .to_dict(). While context snippet 4 shows the Llama converter follows the same pattern (so this is established behavior), it creates a disconnect with the base class contract in converter.py:193-210.

Consider updating the return type to List[Dict[str, Any]] to accurately reflect what is returned, or coordinate a broader fix across all converters.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@modelopt/torch/puzzletron/anymodel/models/qwen3/qwen3_converter.py` at line
33, The function create_block_configs_from_main_config currently declares a
return type of List[BlockConfig] but actually returns a list of dicts (via
.to_dict()); update the signature to return List[Dict[str, Any]] and add the
necessary typing imports (Dict, Any) so the annotation matches the actual return
value, and scan related converters (e.g., the Llama converter and base converter
methods) to ensure callers expect dicts or adjust them consistently.

modelopt/torch/puzzletron/anymodel/models/qwen3/qwen3_model_descriptor.py (1)

64-65: Add return type annotation for consistency with base class.

The base class in model_descriptor.py:42-54 declares block_config_to_layer_overrides with return type Dict[str, Any]. Adding the annotation here would improve type safety and IDE support.

♻️ Proposed fix

+from typing import Any, Dict, List
-from typing import Dict, List
...
     `@staticmethod`
-    def block_config_to_layer_overrides(block_config: BlockConfig):
+    def block_config_to_layer_overrides(block_config: BlockConfig) -> Dict[str, Any]:
         return {

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@modelopt/torch/puzzletron/anymodel/models/qwen3/qwen3_model_descriptor.py`
around lines 64 - 65, The method block_config_to_layer_overrides in
qwen3_model_descriptor.py is missing the return type annotation declared on the
base class; update its signature to include -> Dict[str, Any] (import Dict and
Any from typing if not already present) so it matches the base class declaration
in model_descriptor.py (see block_config_to_layer_overrides and BlockConfig) and
restores proper type safety and IDE support.

modelopt/torch/puzzletron/anymodel/models/qwen2/qwen2_model_descriptor.py (1)

15-15: Narrow the file-wide mypy suppression.

Line 15 (# mypy: ignore-errors) disables type checking for the entire file. Multiple static methods lack return type annotations (e.g., decoder_layer_cls(), input_embedding_name(), layer_block_name(index: int), etc.). Add return type annotations to these methods and use targeted # type: ignore comments only where necessary (e.g., for untyped third-party imports from transformers).

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@modelopt/torch/puzzletron/anymodel/models/qwen2/qwen2_model_descriptor.py` at
line 15, The file-wide "# mypy: ignore-errors" must be removed and replaced by
precise type annotations on the static methods (e.g., add return types to
decoder_layer_cls, input_embedding_name, layer_block_name(index: int), and other
static methods) and only use targeted "# type: ignore" comments for specific
untyped third-party imports; for example declare decoder_layer_cls() ->
Type[Any] or the proper transformer layer class, input_embedding_name() -> str,
layer_block_name(index: int) -> str, and any other methods with explicit return
types so mypy can validate the file while leaving isolated ignores on the import
lines that actually require them.

modelopt/torch/puzzletron/anymodel/models/qwen3_vl/qwen3_vl_converter.py (1)

74-76: Consider using mprint instead of print for logging.

Other modules in this PR use mprint from modelopt.torch.puzzletron.tools.logger for consistent logging. Using print() directly may not integrate well with logging configurations.

♻️ Suggested change

+from modelopt.torch.puzzletron.tools.logger import mprint
+
 ...
 
-        print(
+        mprint(
             f"Created {len(block_configs)} block configs for Qwen3-VL MoE (decoder_sparse_step={decoder_sparse_step})"
         )

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@modelopt/torch/puzzletron/anymodel/models/qwen3_vl/qwen3_vl_converter.py`
around lines 74 - 76, Replace the direct print call with the project logging
helper: import mprint from modelopt.torch.puzzletron.tools.logger and change the
print(...) in qwen3_vl_converter.py to mprint(f"Created {len(block_configs)}
block configs for Qwen3-VL MoE (decoder_sparse_step={decoder_sparse_step})") so
the message uses the consistent logger; ensure the import is added near other
imports and no other behavior changes to block_configs or decoder_sparse_step.

modelopt/torch/puzzletron/subblock_stats/calc_subblock_stats.py (1)

27-27: Setting environment variable at module import time has global side effects.

Setting os.environ["PYTORCH_CUDA_ALLOC_CONF"] at the top level of the module means this side effect occurs whenever any code imports this module, even if that code doesn't use the functions here. Consider moving this inside the functions that need it, or making it conditional.

♻️ Suggested change

-os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
+def _set_cuda_alloc_config():
+    """Set CUDA allocator config for memory-intensive operations."""
+    if "PYTORCH_CUDA_ALLOC_CONF" not in os.environ:
+        os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"

Then call _set_cuda_alloc_config() at the start of calculate_subblock_stats or calculate_subblock_stats_for_puzzle_dir.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@modelopt/torch/puzzletron/subblock_stats/calc_subblock_stats.py` at line 27,
The module currently sets os.environ["PYTORCH_CUDA_ALLOC_CONF"] at import time
which creates a global side effect; instead create a small helper (e.g.
_set_cuda_alloc_config()) that sets that env var only when needed, remove the
top-level assignment, and call _set_cuda_alloc_config() at the start of
calculate_subblock_stats and calculate_subblock_stats_for_puzzle_dir so the CUDA
config is applied only when those functions run (ensure the helper is idempotent
to avoid repeated sets).

modelopt/torch/puzzletron/tools/bypassed_training/init_child_from_parent.py (2)

165-166: Remove duplicate variable assignment.

actual_io_workers is computed on line 165 and then recomputed identically on line 181. The second assignment is redundant.

♻️ Suggested fix

     # Profile _save_checkpoint with automatic I/O worker calculation
     mprint("Starting _save_checkpoint...")
-    actual_io_workers = max_workers if max_workers else "auto"
-    mprint(f"I/O Settings: max_workers={actual_io_workers}")
+    mprint(f"I/O Settings: max_workers={max_workers if max_workers else 'auto'}")
     start_time = time.time()
     ...
     # Print profiling summary with actual worker counts used
     total_core_time = create_child_state_dict_time + save_checkpoint_time
     actual_layer_workers = max_layer_workers if max_layer_workers else "auto"
-    actual_io_workers = max_workers if max_workers else "auto"
+    actual_io_workers = max_workers if max_workers else "auto"  # keep this one

Also applies to: 180-181

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@modelopt/torch/puzzletron/tools/bypassed_training/init_child_from_parent.py`
around lines 165 - 166, Remove the duplicated assignment of actual_io_workers
and its associated mprint call in init_child_from_parent.py: keep the first
computation (actual_io_workers = max_workers if max_workers else "auto") and the
initial mprint("I/O Settings:..."), and delete the later identical reassignment
and print (the repeated lines around where actual_io_workers is set again near
line 181). This removes the redundant code while leaving any subsequent uses of
actual_io_workers and the mprint calls (function mprint) intact.

---

90-91: Add error handling for JSON parsing.

If model_config_overrides_dict is a malformed JSON string, json.loads() will raise a JSONDecodeError. Consider wrapping this in a try-except with a more informative error message.

🛡️ Suggested improvement

     # Parse JSON if string
     if isinstance(model_config_overrides_dict, str):
-        model_config_overrides_dict = json.loads(model_config_overrides_dict)
+        try:
+            model_config_overrides_dict = json.loads(model_config_overrides_dict)
+        except json.JSONDecodeError as e:
+            raise ValueError(f"Invalid JSON in model_config_overrides_dict: {e}") from e

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@modelopt/torch/puzzletron/tools/bypassed_training/init_child_from_parent.py`
around lines 90 - 91, Wrap the json.loads call that parses
model_config_overrides_dict in a try-except that catches json.JSONDecodeError,
so malformed JSON yields a clear, actionable error; specifically, around the
block that checks isinstance(model_config_overrides_dict, str) and calls
json.loads, catch JSONDecodeError and raise a ValueError (or log and raise) with
a message that includes that model_config_overrides_dict failed to parse and
include the offending string (or truncated preview) and the original exception
message to aid debugging.

modelopt/torch/puzzletron/anymodel/models/nemotron_h/nemotron_h_model_descriptor.py (2)

119-126: Consider adding defensive checks for missing attributes.

_block_no_op_post_init accesses decoder_layer.config.block_configs[decoder_layer.layer_idx] directly. If decoder_layer lacks config, block_configs, or layer_idx attributes (e.g., when called on a DummyBlock during initialization), this will raise an AttributeError.

🛡️ Suggested defensive check

 `@staticmethod`
 def _block_no_op_post_init(decoder_layer):
     """
     Due to the subblock structure of NemotronH always one of the subblock is set to no-op, for a real no-op both attention & ffn no-op should be set to True.
     """
+    if not hasattr(decoder_layer, 'config') or not hasattr(decoder_layer, 'layer_idx'):
+        return
     block_config = decoder_layer.config.block_configs[decoder_layer.layer_idx]
     if block_config.ffn.no_op and block_config.attention.no_op:
         decoder_layer.norm = Same()
         decoder_layer.mixer = MatchingZeros()

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In
`@modelopt/torch/puzzletron/anymodel/models/nemotron_h/nemotron_h_model_descriptor.py`
around lines 119 - 126, The _block_no_op_post_init function assumes
decoder_layer has config, block_configs, and layer_idx; add defensive checks at
the start of _block_no_op_post_init to avoid AttributeError (e.g., if called on
a DummyBlock): verify hasattr(decoder_layer, "config"),
hasattr(decoder_layer.config, "block_configs"), and hasattr(decoder_layer,
"layer_idx") (or use a try/except AttributeError), ensure
decoder_layer.layer_idx is an int within range of
decoder_layer.config.block_configs, and return early if any check fails; then
proceed to access block_config and apply the existing ffn/attention no_op logic
on decoder_layer.norm and decoder_layer.mixer.

---

40-52: Dynamic module discovery may be slow for large module trees.

The get_dynamic_modules function walks all packages under transformers_modules and imports each module to find matching classes. This could be slow if the module tree is large. Consider caching the result or limiting the search scope if performance becomes an issue.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In
`@modelopt/torch/puzzletron/anymodel/models/nemotron_h/nemotron_h_model_descriptor.py`
around lines 40 - 52, get_dynamic_modules currently imports every module under
transformers_modules which can be slow; to fix, add caching and an optional
scope parameter: decorate get_dynamic_modules with functools.lru_cache (or
maintain a module-level cache keyed by module_cls_str and optional scope), and
add an optional argument (e.g., package_whitelist: Optional[Iterable[str]] or
max_results/first_match: int) so the caller can limit search to specific
subpackages or stop after N matches; update the function body to skip packages
not in the whitelist and to short-circuit when max_results is reached, keeping
references to get_dynamic_modules, transformers_modules, matches, and the
import/package iteration logic so changes are easy to locate.

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes